#include <mcp/transport/transport.h>
#include <mcp/error.h>
#include <mcp/json/value.h>
#include <iostream>
#include <sstream>
#include <memory>
#include <stdexcept>

namespace mcp {
namespace transport {

// Transport implementation
struct Transport::Impl {
    std::string transport_id;
    TransportState state;
    MessageHandler message_handler;
    ErrorHandler error_handler;
    std::chrono::steady_clock::time_point creation_time;
    mutable std::mutex state_mutex;
    mutable std::mutex handler_mutex;

    Impl(const std::string& id)
        : transport_id(id), state(TransportState::Disconnected),
          creation_time(std::chrono::steady_clock::now()) {}
};

Transport::Transport(const std::string& id)
    : impl_(std::make_unique<Impl>(id)) {}

Transport::~Transport() = default;

std::string Transport::transport_id() const {
    return impl_->transport_id;
}

TransportState Transport::state() const {
    std::lock_guard<std::mutex> lock(impl_->state_mutex);
    return impl_->state;
}

std::chrono::steady_clock::duration Transport::uptime() const {
    return std::chrono::steady_clock::now() - impl_->creation_time;
}

void Transport::set_message_handler(MessageHandler handler) {
    std::lock_guard<std::mutex> lock(impl_->handler_mutex);
    impl_->message_handler = std::move(handler);
}

void Transport::set_error_handler(ErrorHandler handler) {
    std::lock_guard<std::mutex> lock(impl_->handler_mutex);
    impl_->error_handler = std::move(handler);
}

void Transport::set_state(TransportState new_state) {
    std::lock_guard<std::mutex> lock(impl_->state_mutex);
    impl_->state = new_state;
}

void Transport::handle_message(const json::Value& message) {
    std::lock_guard<std::mutex> lock(impl_->handler_mutex);
    if (impl_->message_handler) {
        impl_->message_handler(message);
    }
}

void Transport::handle_error(const std::string& error_message, std::exception_ptr exception) {
    std::lock_guard<std::mutex> lock(impl_->handler_mutex);
    if (impl_->error_handler) {
        impl_->error_handler(error_message, exception);
    } else {
        // Default error handling
        std::cerr << "Transport error (" << impl_->transport_id << "): " << error_message << std::endl;
    }
}

// StdioTransport implementation
struct StdioTransport::Impl {
    std::thread read_thread;
    std::thread write_thread;
    std::queue<json::Value> message_queue;
    std::queue<std::string> write_queue;
    mutable std::mutex queue_mutex;
    mutable std::mutex write_mutex;
    std::condition_variable queue_cv;
    std::condition_variable write_cv;
    std::atomic<bool> running{false};
    std::atomic<bool> should_stop{false};
    std::chrono::steady_clock::time_point last_message_time;

    Impl() : last_message_time(std::chrono::steady_clock::now()) {}
};

StdioTransport::StdioTransport()
    : Transport("stdio"), impl_(std::make_unique<Impl>()) {}

StdioTransport::~StdioTransport() {
    stop();
}

void StdioTransport::start() {
    if (impl_->running) {
        return;
    }

    impl_->running = true;
    impl_->should_stop = false;
    set_state(TransportState::Connecting);

    try {
        // Start read thread
        impl_->read_thread = std::thread([this]() {
            read_loop();
        });

        // Start write thread
        impl_->write_thread = std::thread([this]() {
            write_loop();
        });

        set_state(TransportState::Connected);
    } catch (const std::exception& e) {
        set_state(TransportState::Error);
        handle_error("Failed to start stdio transport", std::current_exception());
    }
}

void StdioTransport::stop() {
    if (!impl_->running) {
        return;
    }

    impl_->should_stop = true;

    // Notify threads to stop
    impl_->queue_cv.notify_all();
    impl_->write_cv.notify_all();

    // Wait for threads to finish
    if (impl_->read_thread.joinable()) {
        impl_->read_thread.join();
    }

    if (impl_->write_thread.joinable()) {
        impl_->write_thread.join();
    }

    impl_->running = false;
    set_state(TransportState::Disconnected);
}

bool StdioTransport::is_connected() const {
    return impl_->running && state() == TransportState::Connected;
}

void StdioTransport::send_message(const json::Value& message) {
    if (!is_connected()) {
        throw std::runtime_error("Transport not connected");
    }

    // Serialize message to JSON string
    std::string json_str = message.serialize() + "\n";

    {
        std::lock_guard<std::mutex> lock(impl_->write_mutex);
        impl_->write_queue.push(std::move(json_str));
    }

    impl_->write_cv.notify_one();
}

std::optional<json::Value> StdioTransport::receive_message(std::chrono::milliseconds timeout) {
    std::unique_lock<std::mutex> lock(impl_->queue_mutex);

    if (impl_->message_queue.empty()) {
        if (timeout.count() > 0) {
            impl_->queue_cv.wait_for(lock, timeout, [this] {
                return !impl_->message_queue.empty() || impl_->should_stop;
            });
        } else {
            impl_->queue_cv.wait(lock, [this] {
                return !impl_->message_queue.empty() || impl_->should_stop;
            });
        }
    }

    if (impl_->message_queue.empty() || impl_->should_stop) {
        return std::nullopt;
    }

    auto message = std::move(impl_->message_queue.front());
    impl_->message_queue.pop();
    impl_->last_message_time = std::chrono::steady_clock::now();

    return message;
}

std::chrono::steady_clock::duration StdioTransport::last_message_age() const {
    return std::chrono::steady_clock::now() - impl_->last_message_time;
}

void StdioTransport::read_loop() {
    try {
        std::string line;
        while (!impl_->should_stop && std::getline(std::cin, line)) {
            if (line.empty()) {
                continue;
            }

            try {
                // Parse JSON message
                auto message = json::Value::parse(line);
                if (!message) {
                    handle_error("Failed to parse JSON message: " + line, nullptr);
                    continue;
                }

                // Add to message queue
                {
                    std::lock_guard<std::mutex> lock(impl_->queue_mutex);
                    impl_->message_queue.push(std::move(*message));
                }

                // Notify waiting threads
                impl_->queue_cv.notify_one();

                // Handle the message
                handle_message(*message);
            } catch (const std::exception& e) {
                handle_error("Error processing message: " + std::string(e.what()), std::current_exception());
            }
        }
    } catch (const std::exception& e) {
        if (!impl_->should_stop) {
            handle_error("Read loop error: " + std::string(e.what()), std::current_exception());
        }
    }
}

void StdioTransport::write_loop() {
    try {
        while (!impl_->should_stop) {
            std::string message;

            {
                std::unique_lock<std::mutex> lock(impl_->write_mutex);
                if (impl_->write_queue.empty()) {
                    impl_->write_cv.wait(lock, [this] {
                        return !impl_->write_queue.empty() || impl_->should_stop;
                    });
                }

                if (impl_->should_stop) {
                    break;
                }

                message = std::move(impl_->write_queue.front());
                impl_->write_queue.pop();
            }

            // Write message to stdout
            std::cout << message << std::flush;
        }
    } catch (const std::exception& e) {
        if (!impl_->should_stop) {
            handle_error("Write loop error: " + std::string(e.what()), std::current_exception());
        }
    }
}

// HTTP transport implementation
struct HttpTransport::Impl {
    std::string base_url;
    std::string session_id;
    std::chrono::steady_clock::time_point last_activity;
    std::queue<json::Value> pending_messages;
    mutable std::mutex queue_mutex;
    mutable std::mutex activity_mutex;
    std::condition_variable queue_cv;

    Impl(std::string url)
        : base_url(std::move(url)), last_activity(std::chrono::steady_clock::now()) {}
};

HttpTransport::HttpTransport(const std::string& base_url)
    : Transport("http"), impl_(std::make_unique<Impl>(base_url)) {}

HttpTransport::~HttpTransport() = default;

void HttpTransport::start() {
    set_state(TransportState::Connected);
    impl_->last_activity = std::chrono::steady_clock::now();
}

void HttpTransport::stop() {
    set_state(TransportState::Disconnected);
}

bool HttpTransport::is_connected() const {
    return state() == TransportState::Connected;
}

void HttpTransport::send_message(const json::Value& message) {
    if (!is_connected()) {
        throw std::runtime_error("HTTP transport not connected");
    }

    {
        std::lock_guard<std::mutex> lock(impl_->queue_mutex);
        impl_->pending_messages.push(message);
    }

    impl_->queue_cv.notify_one();
    update_activity();

    // In a real implementation, this would send the message via HTTP
    // For now, we'll just simulate it
    std::cout << "[HTTP] Would send message to " << impl_->base_url << ": " << message.serialize() << std::endl;
}

std::optional<json::Value> HttpTransport::receive_message(std::chrono::milliseconds timeout) {
    std::unique_lock<std::mutex> lock(impl_->queue_mutex);

    if (impl_->pending_messages.empty()) {
        if (timeout.count() > 0) {
            impl_->queue_cv.wait_for(lock, timeout, [this] {
                return !impl_->pending_messages.empty();
            });
        } else {
            impl_->queue_cv.wait(lock, [this] {
                return !impl_->pending_messages.empty();
            });
        }
    }

    if (impl_->pending_messages.empty()) {
        return std::nullopt;
    }

    // For simulation purposes, create a response message
    json::Value response = json::Value::object({
        {"jsonrpc", "2.0"},
        {"id", 1},
        {"result", json::Value::object({
            {"status", "success"},
            {"message", "HTTP response received"}
        })}
    });

    impl_->pending_messages.pop();
    update_activity();

    return response;
}

void HttpTransport::set_session_id(const std::string& session_id) {
    impl_->session_id = session_id;
}

std::string HttpTransport::session_id() const {
    return impl_->session_id;
}

std::string HttpTransport::base_url() const {
    return impl_->base_url;
}

void HttpTransport::update_activity() {
    std::lock_guard<std::mutex> lock(impl_->activity_mutex);
    impl_->last_activity = std::chrono::steady_clock::now();
}

std::chrono::steady_clock::duration HttpTransport::last_activity_age() const {
    std::lock_guard<std::mutex> lock(impl_->activity_mutex);
    return std::chrono::steady_clock::now() - impl_->last_activity;
}

// WebSocket transport implementation
struct WebSocketTransport::Impl {
    std::string url;
    std::thread connection_thread;
    std::thread message_thread;
    std::queue<json::Value> message_queue;
    std::queue<json::Value> send_queue;
    mutable std::mutex queue_mutex;
    mutable std::mutex send_mutex;
    std::condition_variable queue_cv;
    std::condition_variable send_cv;
    std::atomic<bool> connected{false};
    std::atomic<bool> should_stop{false};
    std::chrono::steady_clock::time_point last_message_time;

    Impl(std::string ws_url) : url(std::move(ws_url)) {}
};

WebSocketTransport::WebSocketTransport(const std::string& url)
    : Transport("websocket"), impl_(std::make_unique<Impl>(url)) {}

WebSocketTransport::~WebSocketTransport() {
    stop();
}

void WebSocketTransport::start() {
    if (impl_->connected) {
        return;
    }

    impl_->should_stop = false;
    set_state(TransportState::Connecting);

    try {
        // Start connection thread
        impl_->connection_thread = std::thread([this]() {
            connection_loop();
        });

        // Start message processing thread
        impl_->message_thread = std::thread([this]() {
            message_loop();
        });

        set_state(TransportState::Connected);
        impl_->connected = true;
    } catch (const std::exception& e) {
        set_state(TransportState::Error);
        handle_error("Failed to start WebSocket transport", std::current_exception());
    }
}

void WebSocketTransport::stop() {
    if (!impl_->connected) {
        return;
    }

    impl_->should_stop = true;
    impl_->connected = false;

    // Notify threads to stop
    impl_->queue_cv.notify_all();
    impl_->send_cv.notify_all();

    // Wait for threads to finish
    if (impl_->connection_thread.joinable()) {
        impl_->connection_thread.join();
    }

    if (impl_->message_thread.joinable()) {
        impl_->message_thread.join();
    }

    set_state(TransportState::Disconnected);
}

bool WebSocketTransport::is_connected() const {
    return impl_->connected && state() == TransportState::Connected;
}

void WebSocketTransport::send_message(const json::Value& message) {
    if (!is_connected()) {
        throw std::runtime_error("WebSocket transport not connected");
    }

    {
        std::lock_guard<std::mutex> lock(impl_->send_mutex);
        impl_->send_queue.push(message);
    }

    impl_->send_cv.notify_one();

    // In a real implementation, this would send via WebSocket
    std::cout << "[WebSocket] Would send message to " << impl_->url << ": " << message.serialize() << std::endl;
}

std::optional<json::Value> WebSocketTransport::receive_message(std::chrono::milliseconds timeout) {
    std::unique_lock<std::mutex> lock(impl_->queue_mutex);

    if (impl_->message_queue.empty()) {
        if (timeout.count() > 0) {
            impl_->queue_cv.wait_for(lock, timeout, [this] {
                return !impl_->message_queue.empty() || impl_->should_stop;
            });
        } else {
            impl_->queue_cv.wait(lock, [this] {
                return !impl_->message_queue.empty() || impl_->should_stop;
            });
        }
    }

    if (impl_->message_queue.empty() || impl_->should_stop) {
        return std::nullopt;
    }

    auto message = std::move(impl_->message_queue.front());
    impl_->message_queue.pop();
    impl_->last_message_time = std::chrono::steady_clock::now();

    return message;
}

std::string WebSocketTransport::url() const {
    return impl_->url;
}

void WebSocketTransport::connection_loop() {
    try {
        // In a real implementation, this would establish WebSocket connection
        std::cout << "[WebSocket] Connecting to " << impl_->url << std::endl;

        // Simulate connection establishment
        std::this_thread::sleep_for(std::chrono::milliseconds(100));

        while (!impl_->should_stop) {
            // Simulate receiving messages
            std::this_thread::sleep_for(std::chrono::milliseconds(1000));

            // Create a simulated incoming message
            json::Value message = json::Value::object({
                {"jsonrpc", "2.0"},
                {"method", "ping"},
                {"params", json::Value::object({})}
            });

            {
                std::lock_guard<std::mutex> lock(impl_->queue_mutex);
                impl_->message_queue.push(std::move(message));
            }

            impl_->queue_cv.notify_one();
            handle_message(message);
        }
    } catch (const std::exception& e) {
        if (!impl_->should_stop) {
            handle_error("WebSocket connection error: " + std::string(e.what()), std::current_exception());
        }
    }
}

void WebSocketTransport::message_loop() {
    try {
        while (!impl_->should_stop) {
            json::Value message;

            {
                std::unique_lock<std::mutex> lock(impl_->send_mutex);
                if (impl_->send_queue.empty()) {
                    impl_->send_cv.wait(lock, [this] {
                        return !impl_->send_queue.empty() || impl_->should_stop;
                    });
                }

                if (impl_->should_stop) {
                    break;
                }

                message = std::move(impl_->send_queue.front());
                impl_->send_queue.pop();
            }

            // In a real implementation, this would send via WebSocket
            // For now, just log it
            std::cout << "[WebSocket] Message sent: " << message.serialize() << std::endl;
        }
    } catch (const std::exception& e) {
        if (!impl_->should_stop) {
            handle_error("WebSocket message loop error: " + std::string(e.what()), std::current_exception());
        }
    }
}

std::chrono::steady_clock::duration WebSocketTransport::last_message_age() const {
    return std::chrono::steady_clock::now() - impl_->last_message_time;
}

// SSE (Server-Sent Events) transport implementation
struct SseTransport::Impl {
    std::string url;
    std::thread receive_thread;
    std::queue<json::Value> message_queue;
    mutable std::mutex queue_mutex;
    std::condition_variable queue_cv;
    std::atomic<bool> connected{false};
    std::atomic<bool> should_stop{false};
    std::chrono::steady_clock::time_point last_message_time;

    Impl(std::string sse_url) : url(std::move(sse_url)) {}
};

SseTransport::SseTransport(const std::string& url)
    : Transport("sse"), impl_(std::make_unique<Impl>(url)) {}

SseTransport::~SseTransport() {
    stop();
}

void SseTransport::start() {
    if (impl_->connected) {
        return;
    }

    impl_->should_stop = false;
    set_state(TransportState::Connecting);

    try {
        // Start receiving thread
        impl_->receive_thread = std::thread([this]() {
            receive_loop();
        });

        set_state(TransportState::Connected);
        impl_->connected = true;
    } catch (const std::exception& e) {
        set_state(TransportState::Error);
        handle_error("Failed to start SSE transport", std::current_exception());
    }
}

void SseTransport::stop() {
    if (!impl_->connected) {
        return;
    }

    impl_->should_stop = true;
    impl_->connected = false;

    // Notify thread to stop
    impl_->queue_cv.notify_all();

    // Wait for thread to finish
    if (impl_->receive_thread.joinable()) {
        impl_->receive_thread.join();
    }

    set_state(TransportState::Disconnected);
}

bool SseTransport::is_connected() const {
    return impl_->connected && state() == TransportState::Connected;
}

void SseTransport::send_message(const json::Value& message) {
    if (!is_connected()) {
        throw std::runtime_error("SSE transport not connected");
    }

    // SSE is typically receive-only, but we can send via separate HTTP requests
    std::cout << "[SSE] Would send HTTP POST to " << impl_->url << ": " << message.serialize() << std::endl;
}

std::optional<json::Value> SseTransport::receive_message(std::chrono::milliseconds timeout) {
    std::unique_lock<std::mutex> lock(impl_->queue_mutex);

    if (impl_->message_queue.empty()) {
        if (timeout.count() > 0) {
            impl_->queue_cv.wait_for(lock, timeout, [this] {
                return !impl_->message_queue.empty() || impl_->should_stop;
            });
        } else {
            impl_->queue_cv.wait(lock, [this] {
                return !impl_->message_queue.empty() || impl_->should_stop;
            });
        }
    }

    if (impl_->message_queue.empty() || impl_->should_stop) {
        return std::nullopt;
    }

    auto message = std::move(impl_->message_queue.front());
    impl_->message_queue.pop();
    impl_->last_message_time = std::chrono::steady_clock::now();

    return message;
}

std::string SseTransport::url() const {
    return impl_->url;
}

void SseTransport::receive_loop() {
    try {
        // In a real implementation, this would connect to SSE endpoint
        std::cout << "[SSE] Connecting to " << impl_->url << std::endl;

        while (!impl_->should_stop) {
            // Simulate receiving SSE events
            std::this_thread::sleep_for(std::chrono::milliseconds(2000));

            // Create a simulated SSE message
            json::Value message = json::Value::object({
                {"jsonrpc", "2.0"},
                {"method", "notification"},
                {"params", json::Value::object({
                    {"event", "update"},
                    {"data", "Server update received"}
                })}
            });

            {
                std::lock_guard<std::mutex> lock(impl_->queue_mutex);
                impl_->message_queue.push(std::move(message));
            }

            impl_->queue_cv.notify_one();
            handle_message(message);
        }
    } catch (const std::exception& e) {
        if (!impl_->should_stop) {
            handle_error("SSE receive error: " + std::string(e.what()), std::current_exception());
        }
    }
}

std::chrono::steady_clock::duration SseTransport::last_message_age() const {
    return std::chrono::steady_clock::now() - impl_->last_message_time;
}

// Utility functions
std::string transport_state_to_string(TransportState state) {
    switch (state) {
        case TransportState::Disconnected:
            return "Disconnected";
        case TransportState::Connecting:
            return "Connecting";
        case TransportState::Connected:
            return "Connected";
        case TransportState::Error:
            return "Error";
        default:
            return "Unknown";
    }
}

// Factory functions
std::unique_ptr<StdioTransport> create_stdio_transport() {
    return std::make_unique<StdioTransport>();
}

std::unique_ptr<HttpTransport> create_http_transport(const std::string& base_url) {
    return std::make_unique<HttpTransport>(base_url);
}

std::unique_ptr<WebSocketTransport> create_websocket_transport(const std::string& url) {
    return std::make_unique<WebSocketTransport>(url);
}

std::unique_ptr<SseTransport> create_sse_transport(const std::string& url) {
    return std::make_unique<SseTransport>(url);
}

} // namespace transport
} // namespace mcp